#include <dbg.h>

#include <cassert>
using namespace std;

class Solution {
public:
    int countKConstraintSubstrings(string s, int k) {
        int n = s.size();
        int res = 0;

        vector<int> cnt(2, 0);

        int left = 0;
        for (int right = 0; right < n; ++right) {
            ++cnt[s[right] - '0'];
            while (cnt[0] > k && cnt[1] > k) {
                --cnt[s[left] - '0'];
                ++left;
            }
            res += (right - left + 1);
        }
        return res;
    }
};

int main() {
    Solution s;

    string s1 = "10101";
    string s2 = "1010101";
    string s3 = "11111";

    int k1 = 1;
    int k2 = 2;
    int k3 = 1;

    assert(s.countKConstraintSubstrings(s1, k1) == 12);
    assert(s.countKConstraintSubstrings(s2, k2) == 25);
    assert(s.countKConstraintSubstrings(s3, k3) == 15);

    return 0;
}